[KDA] Add TLE inference backend for NVIDIA H800 with up to 1.41x speedup - #1075

Sorry for the late response.

My main concern is maintenance cost over time. TileLang is still a moving target: a compiler version bump can turn into compile errors or, worse, silent numerical errors. Meanwhile the Triton baseline for KDA is still evolving quickly, so a parallel TLE implementation risks going stale and drifting out of sync with the reference path.

Before we take on a second backend to maintain, could you share how you see this staying correct over time — e.g. is there test/CI coverage that would catch drift from the Triton baseline, and how tightly does the TLE version track it?

Thanks for the review. I understand that TileLang was mentioned as an example of the long-term maintenance and compiler-upgrade risks involved.

This PR already includes reproducible TLE correctness and performance commands:

Performance:

# TLE
FLA_TLE_KDA=1 FLA_FLASH_KDA=0 FLA_TILELANG=0 \
python -m benchmarks.ops.run \
  --op chunk_kda_inference --modes fwd

# FlashKDA
FLA_TLE_KDA=0 FLA_FLASH_KDA=1 FLA_TILELANG=0 \
python -m benchmarks.ops.run \
  --op chunk_kda_inference --modes fwd

Correctness:

python -m pytest tests/ops/test_kda.py -k 'tle_kda' -v

To continuously catch compilation, numerical, or performance regressions caused by changes in the Triton baseline or FlagTree upgrades, I can integrate these commands into the existing H100 CI and pin it to the validated flagtree==0.6.0. Future FlagTree upgrades would explicitly trigger the same validation through a version update.

Would you prefer me to add dedicated TLE correctness and performance CI coverage in this PR? If this direction sounds appropriate, I will follow up with the detailed CI design and local validation results in my next update.